#pragma once

#include <thrust/transform_scan.h>

template <typename Container1,
          typename Container2,
          typename UnaryOperation,
          typename BinaryOperation>
struct TransformInclusiveScan
{
  Container1 A;
  Container2 B;
  UnaryOperation unary_op;
  BinaryOperation binary_op;

  template <typename Range1, typename Range2>
  TransformInclusiveScan(const Range1& X, const Range2& Y, UnaryOperation unary_op, BinaryOperation binary_op)
    : A(X.begin(), X.end()),
      B(Y.begin(), Y.end()),
      unary_op(unary_op),
      binary_op(binary_op)
  {}

  void operator()(void)
  {
    thrust::transform_inclusive_scan(A.begin(), A.end(), B.begin(), unary_op, binary_op);
  }
  
  void reset(void)
  {
    // nothing to do
  }
};

template <typename Container1,
          typename Container2,
          typename UnaryOperation,
          typename T,
          typename BinaryOperation>
struct TransformExclusiveScan
{
  Container1 A;
  Container2 B;
  T init;
  UnaryOperation unary_op;
  BinaryOperation binary_op;

  template <typename Range1, typename Range2>
  TransformExclusiveScan(const Range1& X, const Range2& Y, UnaryOperation unary_op, T init, BinaryOperation binary_op)
    : A(X.begin(), X.end()),
      B(Y.begin(), Y.end()),
      init(init),
      unary_op(unary_op),
      binary_op(binary_op)
  {}

  void operator()(void)
  {
    thrust::transform_exclusive_scan(A.begin(), A.end(), B.begin(), unary_op, init, binary_op);
  }
  
  void reset(void)
  {
    // nothing to do
  }
};

